import torch
import gradio as gr
from transformers import BertTokenizer, BertForTokenClassification
import os

# 设置正确的模型路径
MODEL_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models/best_model")
TOKENIZER_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models/tokenizer")

# 加载tokenizer
tokenizer = BertTokenizer.from_pretrained(TOKENIZER_PATH)

# 加载模型（使用from_pretrained而不是load_state_dict）
model = BertForTokenClassification.from_pretrained(MODEL_PATH)
model.eval()

# 预测函数
def extract_opinion_target(text):
    # 对输入文本进行tokenize
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    
    # 获取模型预测
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predictions = torch.argmax(logits, dim=2)
    
    # 解码预测结果
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    token_predictions = [prediction.item() for prediction in predictions[0]]
    
    # 提取评价对象
    # 假设标签为：0=O（非实体），1=B-OT（评价对象开始），2=I-OT（评价对象内部）
    opinion_target = []
    current_ot = []
    
    for token, pred in zip(tokens, token_predictions):
        if pred == 1:  # B-OT
            if current_ot:  # 如果已经在收集一个评价对象，结束它
                opinion_target.append(''.join(current_ot))
                current_ot = []
            if token.startswith("##"):
                current_ot.append(token[2:])
            elif token != "[CLS]" and token != "[SEP]" and token != "[PAD]":
                current_ot.append(token)
        elif pred == 2:  # I-OT
            if token.startswith("##"):
                current_ot.append(token[2:])
            elif token != "[CLS]" and token != "[SEP]" and token != "[PAD]":
                current_ot.append(token)
        elif pred == 0 and current_ot:  # O，且已收集了评价对象
            opinion_target.append(''.join(current_ot))
            current_ot = []
    
    # 检查是否还有未添加的评价对象
    if current_ot:
        opinion_target.append(''.join(current_ot))
    
    # 如果没有找到评价对象，返回提示
    if not opinion_target:
        return "未检测到评价对象"
    
    # 返回找到的所有评价对象
    return " | ".join(opinion_target)

# 创建Gradio接口
iface = gr.Interface(
    fn=extract_opinion_target,
    inputs=gr.Textbox(lines=5, placeholder="请输入包含评价对象的文本..."),
    outputs=gr.Textbox(label="提取的评价对象"),
    title="评价对象提取系统",
    description="输入一段包含评价对象的文本，系统将自动提取评价对象。",
    examples=[
        ["《数据结构》，是2011年中国铁道出版社出版的图书，本书结构清晰、语言精练、注重应用，强调系统性和实用性的结合，适合作为高等学校计算机及相关专业的本科教材或参考书，也可作为计算机爱好者的自学参考书。"],
        ["另外，坐车从水绕四门往黄龙洞约10分钟，有一家\"银都土家餐馆\"味道很好，尤其油炸野蜂蛹、石鸡、金鞭鱼非常美味。"],
        ["色拉寺后一座山，爬上去就可以俯瞰拉萨市中心的景色。"]
    ]
)

if __name__ == "__main__":
    iface.launch()